#!/usr/bin/python3

#author: xiaowei

import sys
import time
from datetime import datetime
import requests
from mesos.v1.master import master_pb2
from mesos.v1 import mesos_pb2


HEADERS = {
    'Content-Type': "application/x-protobuf",
    'Accept': "application/x-protobuf",
    'Message-Accept': "application/recordio",
    'cache-control': "no-cache",
}


class MesosEvent(object):

    def __init__(self, mesosurl):
        self.url = mesosurl

    def call_mesos(self):
        print('%s we subscribe the mesos: %s' %
              (datetime.now().isoformat(' '), self.url))
        subscribed = master_pb2.Call()
        subscribed.type = master_pb2.Call.SUBSCRIBE

        payload = subscribed.SerializeToString()

        self.res = requests.request(
            "POST", self.url, data=payload, headers=HEADERS, stream=True, timeout=(3, 120), )

        try:
            self._api_handle()
        finally:
            if self.res is not None:
                self.res.close()

    def _api_handle(self):
        for i in range(10):  # 此处我们只读取10个Record, 然后结束
            record = self._read_record()
            if record is not None:
                print('\n%s get the %d record, length: %s' %
                      (datetime.now().isoformat(' '), i, len(record)))
                event = self._parse_record(record)
                if event is not None:
                    self._handle_event(event)

    def _handle_event(self, e):
        self.show_event(e)

    def _parse_record(self, r):
        e = master_pb2.Event()
        try:
            e.ParseFromString(r)
        except:
            print(r)
            raise

        return e

    def _read_record(self):
        r = self.res.raw

        # now read record first part, we call it head of record

        head = bytearray()
        pos = -1
        while len(head) < 21:
            # a HEARTBEAT record total size is 4 bytes: 1 byte of the 1st part, 1 byte of LF, 2 bytes of the 3rd part
            # so we only read 4 bytes each time, otherwise we will maybe read more than two records
            data = r.read(4)
            if not data:  # no data, maybe we need wait
                time.sleep(0.1)
                continue
            head.extend(data)
            pos = head.find(b'\n')
            if pos > 0:
                break

        if pos > 0:
            record_body_len = int(head[:pos])
        else:
            raise IOError(
                'can not find record second part, the `LF` Character: ' + str(head))

        # now read record third part, we call it body of record

        if record_body_len <= 0 or record_body_len > 1073741824:  # we assume record body less than 1GiB
            raise ValueError('record data, length error: ' + str(head))

        if pos < len(head) - 1:
            body_remain_read = record_body_len - len(head) + pos + 1
        else:
            body_remain_read = record_body_len

        body = r.read(body_remain_read)
        if body_remain_read < record_body_len:
            body = head[pos+1:] + body

        return body

    def show_event(self, e):

        et = master_pb2.Event.Type.Name(e.type)
        print("This is a %s Event" % (et, ))

        if hasattr(self, '_show_event_' + et.lower()):
            getattr(self, '_show_event_' + et.lower())(e)

    def _show_event_subscribed(self, e):
        showtpl = '''The cluster has %s agents, %d frameworks, %s tasks\n'''
        s = e.subscribed.get_state
        print(showtpl % (len(s.get_agents.agents), len(
            s.get_frameworks.frameworks), len(s.get_tasks.tasks)))

    def _show_event_task_added(self, e):
        a = e.task_added
        print('A new task %s was %s ' %
              (a.task.task_id.value, mesos_pb2.TaskState.Name(a.task.state)))

    def _show_event_task_updated(self, e):
        u = e.task_updated
        print('The task %s was %s ' %
              (u.status.task_id.value, mesos_pb2.TaskState.Name(u.state)))


if __name__ == '__main__':
    if len(sys.argv) > 1 and len(sys.argv[1].strip()) > 0:
        mesosurl = sys.argv[1].strip()
    else:
        mesosurl = "http://10.211.55.6:5050/api/v1"
    MesosEvent(mesosurl).call_mesos()
